"""
弹性机制工具
Resilience Mechanisms

提供重试、断路器和限流功能，提高系统稳定性
Provides retry, circuit breaker, and rate limiting functionality to improve system stability
"""

import logging
import asyncio
import time
import threading
from typing import Dict, Any, Optional, Callable, Union, List
from datetime import datetime, timedelta
from enum import Enum
from functools import wraps
import random

from config import IOT_CONFIG

logger = logging.getLogger(__name__)


class CircuitBreakerState(Enum):
    """断路器状态 / Circuit breaker states"""
    CLOSED = "closed"      # 正常状态 / Normal state
    OPEN = "open"          # 熔断状态 / Open state  
    HALF_OPEN = "half_open"  # 半开状态 / Half-open state


class RetryPolicy:
    """
    重试策略
    Retry Policy
    
    定义重试的行为和参数
    Defines retry behavior and parameters
    """
    
    def __init__(self, max_attempts: int = 3, base_delay: float = 1.0, 
                 max_delay: float = 60.0, backoff_factor: float = 2.0,
                 jitter: bool = True):
        """
        初始化重试策略
        Initialize retry policy
        
        Args:
            max_attempts (int): 最大重试次数 / Maximum retry attempts
            base_delay (float): 基础延迟时间(秒) / Base delay time (seconds)
            max_delay (float): 最大延迟时间(秒) / Maximum delay time (seconds)
            backoff_factor (float): 退避因子 / Backoff factor
            jitter (bool): 是否添加随机抖动 / Whether to add random jitter
        """
        self.max_attempts = max_attempts
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.backoff_factor = backoff_factor
        self.jitter = jitter
        
    def calculate_delay(self, attempt: int) -> float:
        """
        计算延迟时间
        Calculate delay time
        
        Args:
            attempt (int): 当前重试次数 / Current attempt number
            
        Returns:
            float: 延迟时间(秒) / Delay time (seconds)
        """
        # 指数退避 / Exponential backoff
        delay = self.base_delay * (self.backoff_factor ** (attempt - 1))
        delay = min(delay, self.max_delay)
        
        # 添加随机抖动避免惊群效应 / Add random jitter to avoid thundering herd
        if self.jitter:
            delay *= (0.5 + random.random() * 0.5)
            
        return delay


class CircuitBreaker:
    """
    断路器
    Circuit Breaker
    
    当服务失败率过高时，暂时停止请求以保护系统
    Temporarily stops requests when service failure rate is too high to protect the system
    """
    
    def __init__(self, failure_threshold: int = 5, recovery_timeout: float = 60.0,
                 expected_exception: Union[Exception, tuple] = Exception):
        """
        初始化断路器
        Initialize circuit breaker
        
        Args:
            failure_threshold (int): 失败阈值 / Failure threshold
            recovery_timeout (float): 恢复超时时间(秒) / Recovery timeout (seconds)
            expected_exception: 预期异常类型 / Expected exception types
        """
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.expected_exception = expected_exception
        
        self.failure_count = 0
        self.last_failure_time = None
        self.state = CircuitBreakerState.CLOSED
        self._lock = threading.Lock()
        
        self.logger = logging.getLogger(f"{self.__class__.__name__}")
    
    def can_execute(self) -> bool:
        """
        检查是否可以执行请求
        Check if request can be executed
        
        Returns:
            bool: 是否可以执行 / Whether can execute
        """
        with self._lock:
            if self.state == CircuitBreakerState.CLOSED:
                return True
            elif self.state == CircuitBreakerState.OPEN:
                # 检查是否到了恢复时间 / Check if it's time to recover
                if (self.last_failure_time and 
                    time.time() - self.last_failure_time >= self.recovery_timeout):
                    self.state = CircuitBreakerState.HALF_OPEN
                    self.logger.info("断路器进入半开状态 / Circuit breaker entering half-open state")
                    return True
                return False
            else:  # HALF_OPEN
                return True
    
    def record_success(self):
        """记录成功执行 / Record successful execution"""
        with self._lock:
            if self.state == CircuitBreakerState.HALF_OPEN:
                self.state = CircuitBreakerState.CLOSED
                self.failure_count = 0
                self.logger.info("断路器恢复正常状态 / Circuit breaker recovered to closed state")
    
    def record_failure(self, exception: Exception):
        """
        记录失败执行
        Record failed execution
        
        Args:
            exception (Exception): 异常对象 / Exception object
        """
        if not isinstance(exception, self.expected_exception):
            return
            
        with self._lock:
            self.failure_count += 1
            self.last_failure_time = time.time()
            
            if (self.state == CircuitBreakerState.CLOSED and 
                self.failure_count >= self.failure_threshold):
                self.state = CircuitBreakerState.OPEN
                self.logger.warning(
                    f"断路器打开 / Circuit breaker opened after {self.failure_count} failures"
                )
            elif self.state == CircuitBreakerState.HALF_OPEN:
                self.state = CircuitBreakerState.OPEN
                self.logger.warning("断路器重新打开 / Circuit breaker reopened")
    
    def get_state_info(self) -> Dict[str, Any]:
        """
        获取断路器状态信息
        Get circuit breaker state information
        
        Returns:
            Dict[str, Any]: 状态信息 / State information
        """
        with self._lock:
            return {
                'state': self.state.value,
                'failure_count': self.failure_count,
                'failure_threshold': self.failure_threshold,
                'last_failure_time': self.last_failure_time,
                'recovery_timeout': self.recovery_timeout
            }


class RateLimiter:
    """
    限流器
    Rate Limiter
    
    控制请求速率，防止系统过载
    Controls request rate to prevent system overload
    """
    
    def __init__(self, max_requests: int, time_window: float = 60.0):
        """
        初始化限流器
        Initialize rate limiter
        
        Args:
            max_requests (int): 时间窗口内最大请求数 / Maximum requests in time window
            time_window (float): 时间窗口(秒) / Time window (seconds)
        """
        self.max_requests = max_requests
        self.time_window = time_window
        self.requests = []
        self._lock = threading.Lock()
    
    def can_proceed(self) -> bool:
        """
        检查是否可以继续请求
        Check if request can proceed
        
        Returns:
            bool: 是否可以继续 / Whether can proceed
        """
        current_time = time.time()
        
        with self._lock:
            # 清理过期的请求记录 / Clean up expired request records
            cutoff_time = current_time - self.time_window
            self.requests = [req_time for req_time in self.requests if req_time > cutoff_time]
            
            # 检查是否超过限制 / Check if limit is exceeded
            if len(self.requests) >= self.max_requests:
                return False
            
            # 记录当前请求 / Record current request
            self.requests.append(current_time)
            return True
    
    def get_status(self) -> Dict[str, Any]:
        """
        获取限流器状态
        Get rate limiter status
        
        Returns:
            Dict[str, Any]: 状态信息 / Status information
        """
        current_time = time.time()
        with self._lock:
            cutoff_time = current_time - self.time_window
            active_requests = [req_time for req_time in self.requests if req_time > cutoff_time]
            
            return {
                'max_requests': self.max_requests,
                'time_window': self.time_window,
                'current_requests': len(active_requests),
                'remaining_capacity': self.max_requests - len(active_requests)
            }


class ResilienceManager:
    """
    弹性管理器
    Resilience Manager
    
    统一管理重试、断路器和限流功能
    Unified management of retry, circuit breaker, and rate limiting functionality
    """
    
    def __init__(self):
        """初始化弹性管理器 / Initialize resilience manager"""
        self.logger = logging.getLogger(f"{self.__class__.__name__}")
        
        # 从配置获取默认参数 / Get default parameters from config
        error_config = IOT_CONFIG.get('error_handling', {})
        self.default_retry_policy = RetryPolicy(
            max_attempts=error_config.get('max_retries', 3),
            base_delay=error_config.get('retry_delay', 1.0)
        )
        
        # 断路器和限流器实例 / Circuit breaker and rate limiter instances
        self.circuit_breakers: Dict[str, CircuitBreaker] = {}
        self.rate_limiters: Dict[str, RateLimiter] = {}
        self._lock = threading.Lock()
    
    def get_circuit_breaker(self, name: str, failure_threshold: int = None, 
                          recovery_timeout: float = None) -> CircuitBreaker:
        """
        获取或创建断路器
        Get or create circuit breaker
        
        Args:
            name (str): 断路器名称 / Circuit breaker name
            failure_threshold (int): 失败阈值 / Failure threshold
            recovery_timeout (float): 恢复超时 / Recovery timeout
            
        Returns:
            CircuitBreaker: 断路器实例 / Circuit breaker instance
        """
        with self._lock:
            if name not in self.circuit_breakers:
                error_config = IOT_CONFIG.get('error_handling', {})
                threshold = failure_threshold or error_config.get('circuit_breaker_threshold', 5)
                timeout = recovery_timeout or error_config.get('circuit_breaker_timeout', 60)
                
                self.circuit_breakers[name] = CircuitBreaker(
                    failure_threshold=threshold,
                    recovery_timeout=timeout
                )
                self.logger.info(f"创建断路器 / Created circuit breaker: {name}")
            
            return self.circuit_breakers[name]
    
    def get_rate_limiter(self, name: str, max_requests: int = None, 
                        time_window: float = None) -> RateLimiter:
        """
        获取或创建限流器
        Get or create rate limiter
        
        Args:
            name (str): 限流器名称 / Rate limiter name
            max_requests (int): 最大请求数 / Maximum requests
            time_window (float): 时间窗口 / Time window
            
        Returns:
            RateLimiter: 限流器实例 / Rate limiter instance
        """
        with self._lock:
            if name not in self.rate_limiters:
                requests = max_requests or 100
                window = time_window or 60.0
                
                self.rate_limiters[name] = RateLimiter(
                    max_requests=requests,
                    time_window=window
                )
                self.logger.info(f"创建限流器 / Created rate limiter: {name}")
            
            return self.rate_limiters[name]
    
    async def execute_with_retry(self, func: Callable, *args, 
                               retry_policy: RetryPolicy = None, 
                               circuit_breaker_name: str = None,
                               rate_limiter_name: str = None,
                               **kwargs) -> Any:
        """
        带重试机制执行函数
        Execute function with retry mechanism
        
        Args:
            func (Callable): 要执行的函数 / Function to execute
            *args: 函数参数 / Function arguments
            retry_policy (RetryPolicy): 重试策略 / Retry policy
            circuit_breaker_name (str): 断路器名称 / Circuit breaker name
            rate_limiter_name (str): 限流器名称 / Rate limiter name
            **kwargs: 函数关键字参数 / Function keyword arguments
            
        Returns:
            Any: 函数执行结果 / Function execution result
            
        Raises:
            Exception: 最后一次执行的异常 / Exception from last execution
        """
        policy = retry_policy or self.default_retry_policy
        circuit_breaker = None
        rate_limiter = None
        
        # 获取断路器和限流器 / Get circuit breaker and rate limiter
        if circuit_breaker_name:
            circuit_breaker = self.get_circuit_breaker(circuit_breaker_name)
        if rate_limiter_name:
            rate_limiter = self.get_rate_limiter(rate_limiter_name)
        
        last_exception = None
        
        for attempt in range(1, policy.max_attempts + 1):
            try:
                # 检查断路器状态 / Check circuit breaker state
                if circuit_breaker and not circuit_breaker.can_execute():
                    raise Exception("断路器开启，请求被拒绝 / Circuit breaker open, request rejected")
                
                # 检查限流器 / Check rate limiter
                if rate_limiter and not rate_limiter.can_proceed():
                    raise Exception("超过速率限制 / Rate limit exceeded")
                
                # 执行函数 / Execute function
                if asyncio.iscoroutinefunction(func):
                    result = await func(*args, **kwargs)
                else:
                    result = func(*args, **kwargs)
                
                # 记录成功 / Record success
                if circuit_breaker:
                    circuit_breaker.record_success()
                
                self.logger.debug(f"函数执行成功 / Function executed successfully on attempt {attempt}")
                return result
                
            except Exception as e:
                last_exception = e
                
                # 记录失败 / Record failure
                if circuit_breaker:
                    circuit_breaker.record_failure(e)
                
                self.logger.warning(
                    f"函数执行失败 / Function execution failed on attempt {attempt}/{policy.max_attempts}: {str(e)}"
                )
                
                # 如果不是最后一次尝试，等待后重试 / If not last attempt, wait and retry
                if attempt < policy.max_attempts:
                    delay = policy.calculate_delay(attempt)
                    self.logger.info(f"等待 {delay:.2f} 秒后重试 / Waiting {delay:.2f}s before retry")
                    await asyncio.sleep(delay)
                else:
                    self.logger.error(f"所有重试尝试失败 / All retry attempts failed")
        
        # 抛出最后一次的异常 / Raise the last exception
        raise last_exception
    
    def get_all_status(self) -> Dict[str, Any]:
        """
        获取所有弹性组件的状态
        Get status of all resilience components
        
        Returns:
            Dict[str, Any]: 状态信息 / Status information
        """
        with self._lock:
            circuit_breaker_status = {}
            for name, cb in self.circuit_breakers.items():
                circuit_breaker_status[name] = cb.get_state_info()
            
            rate_limiter_status = {}
            for name, rl in self.rate_limiters.items():
                rate_limiter_status[name] = rl.get_status()
            
            return {
                'circuit_breakers': circuit_breaker_status,
                'rate_limiters': rate_limiter_status,
                'default_retry_policy': {
                    'max_attempts': self.default_retry_policy.max_attempts,
                    'base_delay': self.default_retry_policy.base_delay,
                    'max_delay': self.default_retry_policy.max_delay,
                    'backoff_factor': self.default_retry_policy.backoff_factor,
                    'jitter': self.default_retry_policy.jitter
                }
            }
    
    async def cleanup(self):
        """清理资源 / Cleanup resources"""
        with self._lock:
            self.circuit_breakers.clear()
            self.rate_limiters.clear()
        self.logger.info("弹性管理器已清理 / Resilience manager cleaned up")


# 装饰器函数 / Decorator functions

def retry(max_attempts: int = None, base_delay: float = None, 
          circuit_breaker: str = None, rate_limiter: str = None):
    """
    重试装饰器
    Retry decorator
    
    Args:
        max_attempts (int): 最大重试次数 / Maximum retry attempts
        base_delay (float): 基础延迟时间 / Base delay time
        circuit_breaker (str): 断路器名称 / Circuit breaker name
        rate_limiter (str): 限流器名称 / Rate limiter name
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            error_config = IOT_CONFIG.get('error_handling', {})
            retry_policy = RetryPolicy(
                max_attempts=max_attempts or error_config.get('max_retries', 3),
                base_delay=base_delay or error_config.get('retry_delay', 1.0)
            )
            
            return await resilience_manager.execute_with_retry(
                func, *args,
                retry_policy=retry_policy,
                circuit_breaker_name=circuit_breaker,
                rate_limiter_name=rate_limiter,
                **kwargs
            )
        return wrapper
    return decorator


def circuit_breaker(name: str, failure_threshold: int = None, recovery_timeout: float = None):
    """
    断路器装饰器
    Circuit breaker decorator
    
    Args:
        name (str): 断路器名称 / Circuit breaker name
        failure_threshold (int): 失败阈值 / Failure threshold
        recovery_timeout (float): 恢复超时 / Recovery timeout
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            cb = resilience_manager.get_circuit_breaker(name, failure_threshold, recovery_timeout)
            
            if not cb.can_execute():
                raise Exception(f"断路器 {name} 开启，请求被拒绝 / Circuit breaker {name} open, request rejected")
            
            try:
                if asyncio.iscoroutinefunction(func):
                    result = await func(*args, **kwargs)
                else:
                    result = func(*args, **kwargs)
                cb.record_success()
                return result
            except Exception as e:
                cb.record_failure(e)
                raise
        return wrapper
    return decorator


def rate_limit(name: str, max_requests: int = None, time_window: float = None):
    """
    限流装饰器
    Rate limit decorator
    
    Args:
        name (str): 限流器名称 / Rate limiter name
        max_requests (int): 最大请求数 / Maximum requests
        time_window (float): 时间窗口 / Time window
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            rl = resilience_manager.get_rate_limiter(name, max_requests, time_window)
            
            if not rl.can_proceed():
                raise Exception(f"超过速率限制 {name} / Rate limit {name} exceeded")
            
            if asyncio.iscoroutinefunction(func):
                return await func(*args, **kwargs)
            else:
                return func(*args, **kwargs)
        return wrapper
    return decorator


# 全局弹性管理器实例 / Global resilience manager instance
resilience_manager = ResilienceManager()


# 便捷函数 / Convenience functions

async def safe_execute(func: Callable, *args, max_attempts: int = 3, 
                      circuit_breaker_name: str = None, **kwargs) -> Any:
    """
    安全执行函数的便捷方法
    Convenience method to safely execute functions
    
    Args:
        func (Callable): 要执行的函数 / Function to execute
        *args: 函数参数 / Function arguments
        max_attempts (int): 最大重试次数 / Maximum retry attempts
        circuit_breaker_name (str): 断路器名称 / Circuit breaker name
        **kwargs: 函数关键字参数 / Function keyword arguments
        
    Returns:
        Any: 函数执行结果 / Function execution result
    """
    retry_policy = RetryPolicy(max_attempts=max_attempts)
    return await resilience_manager.execute_with_retry(
        func, *args,
        retry_policy=retry_policy,
        circuit_breaker_name=circuit_breaker_name,
        **kwargs
    )


def get_resilience_status() -> Dict[str, Any]:
    """
    获取弹性组件状态的便捷方法
    Convenience method to get resilience component status
    
    Returns:
        Dict[str, Any]: 状态信息 / Status information
    """
    return resilience_manager.get_all_status()


if __name__ == "__main__":
    # 测试代码 / Test code
    async def test_function():
        print("测试函数执行 / Test function executed")
        return "success"
    
    async def failing_function():
        raise Exception("测试异常 / Test exception")
    
    async def main():
        # 测试重试机制 / Test retry mechanism
        try:
            result = await safe_execute(test_function)
            print(f"成功执行 / Successful execution: {result}")
        except Exception as e:
            print(f"执行失败 / Execution failed: {e}")
        
        # 测试断路器 / Test circuit breaker
        try:
            result = await safe_execute(failing_function, max_attempts=2, circuit_breaker_name="test_cb")
        except Exception as e:
            print(f"预期失败 / Expected failure: {e}")
        
        # 查看状态 / Check status
        status = get_resilience_status()
        print(f"弹性组件状态 / Resilience component status: {status}")
    
    asyncio.run(main())